import os,argparse
import numpy as np
from PIL import Image
from models import *
import torch
import torch.nn as nn
import torchvision.transforms as tfs 
import torchvision.utils as vutils
import matplotlib.pyplot as plt
from torchvision.utils import make_grid
from tqdm import tqdm
from torch.utils.data import DataLoader
from data_utils import RESIDE_Dataset
from metrics import ssim, psnr
from models.DFFA import DFFA
abs=os.getcwd()+'/'

parser=argparse.ArgumentParser()
parser.add_argument('--task',type=str,default='its',help='its or ots')
parser.add_argument('--test_imgs',type=str,default='test_imgs',help='Test imgs folder')
parser.add_argument('--device',type=str,default='0,1')
parser.add_argument('--gps',type=int,default=3)
parser.add_argument('--blocks',type=int,default=10)
opt=parser.parse_args()
dataset=opt.task
img_dir=''
model_dir=abs+f'trained_models/DFFANet.pt'
device='cuda' if torch.cuda.is_available() else 'cpu'
ckp=torch.load(model_dir,map_location=device)
net=DFFA(gps=opt.gps,blocks=opt.blocks)
net=nn.DataParallel(net,device_ids=[0])
net.load_state_dict(ckp['model'])

ssims = []
psnrs = []
ITS_test_loader=DataLoader(dataset=RESIDE_Dataset(img_dir,train=False,size='whole img'),batch_size=1,shuffle=False)
loader_test=ITS_test_loader


with torch.no_grad():
    net.eval()
    torch.cuda.empty_cache()
    for i ,(inputs,targets) in tqdm(enumerate(loader_test)):
        inputs=inputs.to(device);targets=targets.to(device)
        pred=net(inputs)
        ssim1=ssim(pred,targets).item()
        psnr1=psnr(pred.detach(),targets.detach())
        ssims.append(ssim1)
        psnrs.append(psnr1)
print(f'mean ssims:{np.mean(ssims)}\t mean psnr:{np.mean(psnrs)}')


